package com.briup.oagg.model.bean.op;

import com.briup.oagg.model.bean.basic.*;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableMap;
import lombok.Data;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 使用建造者模式创建训练配置信息
 *
 * @Author briup
 */
@Data
public class TrainTaskConfig {
    private String taskid; //设置训练id，将来根据这个id将结果和配置关联
    private String task;//训练任务名称 CLASSIFICATION（分类） CLUSTER（聚类） REGRESSION（回归）
    private String[] columns;//计算指标
    private String target; //目标列
    private int n_clusters;//预期聚类个数
    private List<Map<String,Object>> data;//数据集数据信息
    private TaskFlow flow;//算子信息

    private TrainTaskConfig(Builder builder){
        this.taskid = builder.taskid;
        this.task = builder.task;
        this.columns = builder.columns;
        this.target = builder.target;
        this.n_clusters = builder.n_clusters;
        this.data = builder.data;
        this.flow = builder.flow;
    }

    public static final class Builder{
        private String taskid; //设置训练id，将来根据这个id将结果和配置关联
        private String task;// 训练模型类型
        private String[] columns;//计算指标
        private String target; //目标列
        private int n_clusters;//预期聚类个数
        private List<Map<String,Object>> data;//数据集信息
        private TaskFlow flow = new TaskFlow();

        /**
         * 设置训练模型类型
         * @param modelType
         * @return
         */
        public Builder task(String modelType){
            this.task = modelType;
            return this;
        }

        /**
         * 设置分类训练基本配置信息
         * @param modelTrainConfig
         * @return
         */

        /**
         * 设置回顾训练基本配置信息
         * @param modelTrainConfig
         * @return
         */

        public Builder modelTrrainConfig(ModelTrainConfig modelTrainConfig){
            this.taskid = modelTrainConfig.getTrainId();
            this.task = modelTrainConfig.getModelType();
            this.columns = modelTrainConfig.getTrainIndicators().split(",");
            this.target = modelTrainConfig.getTarget();
            return this;
        }

        /**
         * 设置预处理算子信息
         * @param preprocessOperator
         * @return
         */
        public Builder preprocessOperator(PreprocessOperator preprocessOperator){
            TaskFlow.PreProcessFlow pf = new TaskFlow.PreProcessFlow();
            pf.setName(preprocessOperator.getProcessName());
            pf.setConfig(ImmutableMap.of("class_name",preprocessOperator.getProcessClass()));
            this.flow.setPreprocess(Arrays.asList(pf));
            return this;
        }

        /**
         * 设置数据集信息
         * @param dataList
         * @return
         */
        public Builder data(List<Map<String,Object>> dataList){
            this.data = dataList;
            return this;
        }

        /**
         * 设置机器学习算子信息
         * @param learnAlgorithm
         * @return
         */
        public Builder learnAlgorithm(LearnAlgorithm learnAlgorithm){
            MachineLearnFlow ml = new MachineLearnFlow();
            ml.setName(learnAlgorithm.getAlgorithmName());
            ml.setPlatform(learnAlgorithm.getPlatform());
            MachineLearnFlow.MachineLearnOperator op = new MachineLearnFlow.MachineLearnOperator();
            op.setClass_name(learnAlgorithm.getClassName());
            MachineLearnFlow.OperatorConfig opc = new MachineLearnFlow.OperatorConfig();
            opc.setName(learnAlgorithm.getCodeName());
            ObjectMapper mapper = new ObjectMapper();
            try {
                List<MachineLearnFlow.OperatorLayer> operatorLayers = mapper.readValue(learnAlgorithm.getLayers(), new TypeReference<List<MachineLearnFlow.OperatorLayer>>() {});
                opc.setLayers(operatorLayers);
            } catch (JsonProcessingException e) {
                e.printStackTrace();
            }
            op.setConfig(opc);
            ml.setOperator(op);
            this.flow.setMl(Arrays.asList(ml));
            return this;
        }
        public TrainTaskConfig build(){
            return new TrainTaskConfig(this);
        }
    }

    @Data
    public static class ModelTrainingConfig {
        private String trainId;
        private String trainIndicators;
        private String target;
    }
}
